#ifndef AICC_MEM_TO_REG_H
#define AICC_MEM_TO_REG_H

#include "llvm/IR/Function.h"
#include "llvm/IR/PassManager.h"
#include "llvm/Pass.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/IR/Dominators.h"

#include "ai_passes.hpp"

using llvm::dyn_cast;
using llvm::isa;

llvm::PreservedAnalyses
AIMem2Reg::run(llvm::Function &F,
		        llvm::FunctionAnalysisManager &AM)
{
    auto& DT = AM.getResult<llvm::DominatorTreeAnalysis>(F);

    bool changed = false;
	for (auto &BB : F)
	{
		for (auto &I : BB)
		{
            if (auto *LI = dyn_cast<llvm::LoadInst>(&I)) {
                // Find the alloca instruction that defines the loaded value
                auto *AI = dyn_cast<llvm::AllocaInst>(LI->getPointerOperand());
                if (!AI || AI->getParent() != &BB) {
                  continue;
                }

                // Check if the alloca is only stored to once and within the same block
                if (!AI->hasOneUse() || !isa<llvm::StoreInst>(*AI->user_begin())) {
                  continue;
                }

                auto *SI = dyn_cast<llvm::StoreInst>(*AI->user_begin());
                if (SI->getParent() != &BB) {
                  continue;
                }

                // Check dominance relationship between store and load
                if (!DT.dominates(SI, LI)) {
                  continue;
                }

                // Replace all uses of the load with the stored value
                LI->replaceAllUsesWith(SI->getValueOperand());
                LI->eraseFromParent();

                changed = true;
                /**
                 * Давайте еще раз рассмотрим, где нам нужно разместить phi-узел.

                Для существующего load, если последний store может быть в двух (или более) разных BB, тогда мы должны вставить phi-узел для него.

                Тогда проблема заключается в том, как определить, может ли "последний store находиться в двух (или более) разных базовых блоках".
                Тут нужен анализ CFG.

                Узел phi должен быть размещен для этого store в некотором базовом блоке, который содержит load по тому же адресу, когда:
                К лоаду существует путь через данный BB и еще один путь через другой BB.

                Такой базовый блок является в точности одним из элементов dominance базового блока, в котором находится store.
                 */
                
                /**
                 *
                 *
                 * 


                 */
            }
		}
	}

	return llvm::PreservedAnalyses::all();
}
struct MyMem2Reg : PassInfoMixin<MyMem2Reg> {

  PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM) {
    DominatorTree &DT = FAM.getResult<DominatorTreeAnalysis>(F);

    bool Changed = false;
    SmallPtrSet<AllocaInst*, 8> PromotableAllocas;

    // Find promotable allocas
    for (BasicBlock &BB : F) {
      for (Instruction &I : BB) {
        if (auto *AI = dyn_cast<AllocaInst>(&I)) {
          if (isAllocaPromotable(AI)) {
            PromotableAllocas.insert(AI);
          }
        }
      }
    }

    // Process each promotable alloca
    for (AllocaInst *AI : PromotableAllocas) {
      SmallVector<LoadInst*, 8> Loads;
      SmallVector<std::pair<BasicBlock*, Value*>, 8> PredValues;

      // Collect loads and calculate value for each predecessor
      for (User *U : AI->users()) {
        if (auto *LI = dyn_cast<LoadInst>(U)) {
          Loads.push_back(LI);
          for (auto Pred : predecessors(LI->getParent())) {
            PredValues.emplace_back(Pred, findValueDominatingLoad(LI, Pred, DT));
          }
        }
      }

      // Create phi nodes and replace loads
      if (!Loads.empty()) {
        auto *NewVal = new PHINode(AI->getAllocatedType(), PredValues.size(),
                                  AI->getName(), &F.getEntryBlock().front());
        for (auto &Pair : PredValues) {
          NewVal->addIncoming(Pair.second, Pair.first);
        }

        for (auto *LI : Loads) {
          LI->replaceAllUsesWith(NewVal);
          LI->eraseFromParent();
          Changed = true;
        }
      }

      // Erase alloca and any remaining dead instructions
      eraseDeadInstList(AI);
    }

    return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all();
  }

private:
  bool isAllocaPromotable(AllocaInst *AI) {
    // Implement your alloca promotion criteria here, similar to the checks in the simple example.
    // You can consider factors like alloca usage, dominance relationships, and loop dependencies.
    return AI->isStaticAlloca() && AI->hasOneUse(); // Example criteria
  }

  Value *findValueDominatingLoad(LoadInst *LI, BasicBlock *PredBB, DominatorTree &DT) {
    // Implement logic to find the value that dominates the load in the predecessor block.
    // This might involve finding the appropriate store instruction or another value based on dominance.
    return UndefValue::get(LI->getType()); // Placeholder, replace with actual implementation
  }
};
#endif
